Issue 56: Add pathfinder init #117

SamuelBrand1 · 2024-03-06T15:27:54Z

This is a new draft PR aimed at addressing #56 .

The main innovation for this PR occurs here:

Rt-without-renewal/EpiAware/docs/src/examples/getting_started_pf.jl

Lines 245 to 291 in ac788d0

    
           function _epi_aware(y_t, 
        
                   time_steps; 
        
                   epi_model::AbstractEpiModel, 
        
                   latent_model::AbstractLatentModel, 
        
                   observation_model::AbstractObservationModel, 
        
                   nsamples, 
        
                   nchains, 
        
                   pf_ndraws = 50, 
        
                   pf_nruns = 4, 
        
                   fixed_parameters = (;), 
        
                   pos_shift = 1e-6, 
        
                   executor = Transducers.ThreadedEx(), 
        
                   adtype = AutoReverseDiff(true), 
        
                   maxiters = 10, 
        
                   kwargs...) 
        
               gen_mdl = make_epi_aware(missing, time_horizon; epi_model, 
        
                   latent_model, observation_model, pos_shift) |> 
        
                         mdl -> fix(mdl, fixed_parameters) 
        
               mdl = make_epi_aware(y_t, time_horizon; epi_model, 
        
                   latent_model, observation_model, pos_shift) |> 
        
                     mdl -> fix(mdl, fixed_parameters) 
        
               safe_mdl = make_epi_aware(y_t, time_horizon, Val(:safe_mode); epi_model, 
        
                   latent_model, observation_model, pos_shift) |> 
        
                          mdl -> fix(mdl, fixed_parameters) 
        
               mpf = multipathfinder(safe_mdl, max(pf_ndraws, nchains); 
        
                   nruns = pf_nruns, 
        
                   executor, 
        
                   maxiters, 
        
                   kwargs...) 
        
               init_params = collect.(eachrow(mpf.draws_transformed.value[1:nchains, :, 1])) 
        
               chn = sample(mdl, 
        
                   NUTS(; adtype), 
        
                   MCMCThreads(), 
        
                   nsamples ÷ nchains, 
        
                   nchains; 
        
                   init_params = init_params, 
        
                   drop_warmup = true) 
        
               return chn, (; pathfinder_res = mpf, 
        
                   inference_mdl = mdl, 
        
                   generative_mdl = gen_mdl) 
        
           end

The aim is to ultimately have _epi_aware as an _internal_method as part of a wider workflow. But for this PR it is getting this functionality available.

The current first pass behaviour is for _epi_aware to take in the information required to act on make_epi_aware; then do multipathfinder and pass the initial parameters to sample (as per this workflow).

Close #56

… then render into docs

…ut-renewal into add-doc-examples

EpiAware/docs/src/examples/getting_started.jl

SamuelBrand1 · 2024-03-07T11:39:08Z

I took a step back and rather than trying to implement a full loop of model definition -> init -> NUTS, I've refocussed on having a robust initialisation method.

So now I implement an exported manypathfinder function; this is similar to multipathfinder but aimed at finding a good place in parameter space to initialise NUTS in a situation where early large sample values can cause errors.

manypathfinder does nruns threaded parallel runs of pathfinder and selects the best ELBO estimate run with fails as -Inf.

Co-authored-by: Sam Abbott <[email protected]>

seabbs

This is much nicer I think. Some safety and quality of life comments + the obvious big merge issue.

EpiAware/docs/src/examples/getting_started.jl

EpiAware/src/inference-methods.jl

EpiAware/test/test_inference-methods.jl

seabbs · 2024-03-07T11:50:50Z

Once this has stabilised and we have taken it out for a spin it would be interesting to float this wrapper to Pathfinder.jl

…thout-renewal into add-pathfinder-init

codecov-commenter · 2024-03-07T12:13:02Z

Codecov Report

Attention: Patch coverage is 97.05882% with 1 lines in your changes are missing coverage. Please review.

Project coverage is 97.61%. Comparing base (6fd57d8) to head (6deb30a).

Files	Patch %	Lines
EpiAware/src/inference-methods.jl	97.05%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #117      +/-   ##
==========================================
- Coverage   97.76%   97.61%   -0.15%     
==========================================
  Files           6        7       +1     
  Lines         134      168      +34     
==========================================
+ Hits          131      164      +33     
- Misses          3        4       +1

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

SamuelBrand1 · 2024-03-07T12:59:22Z

I've increased the testing for manypathfinder:

More expected failures.
Verbose fail if retries all fail.
KS test that the best pathfinder run is a good approximation to a simple normal (if it isn't then thats an obvious problem).

seabbs

We discussed face to face reordering the inputs to match multipathfinder but think that has yet to happen? Once that is done and we have a new issue for falling back on more functionality from multipathfinder I think this is good to go. It also sounds like we have some confusion about the aim with the docstrings but again perhaps that needs a new issue to discuss..

SamuelBrand1 · 2024-03-08T11:51:44Z

We discussed face to face reordering the inputs to match multipathfinder but think that has yet to happen? Once that is done and we have a new issue for falling back on more functionality from multipathfinder I think this is good to go. It also sounds like we have some confusion about the aim with the docstrings but again perhaps that needs a new issue to discuss..

I've changed the input to match manypathfinder input to multipathfinder, e.g. manypathfinder(mdl, ndraws; kwargs...) -> PathfinderResult is the current signature.

SamuelBrand1 · 2024-03-08T11:52:23Z

After a quick scoping, I don't have immediate plans to use more multipathfinder functionality under the hood anymore.

seabbs · 2024-03-08T12:29:46Z

After a quick scoping, I don't have immediate plans to use more multipathfinder functionality under the hood anymore.

Why

seabbs

LGTM

SamuelBrand1 · 2024-03-08T12:33:57Z

After a quick scoping, I don't have immediate plans to use more multipathfinder functionality under the hood anymore.

Why

I thought that the try catch pattern in manypathfinder could be replaced by fail_on_nonfinite = false in pathfinder kwargs, but the behaviour of the Distributions.NegativeBinomial distribution constructor chucks an error if we find a NaN during warm-up. manypathfinder is resilient to that fail mode.

SamuelBrand1 and others added 30 commits March 1, 2024 18:21

update

0b7c8e1

added doc deps

d86b12d

started getting started notebook

dfeaf54

Update getting_started.jl

2ce57b4

Getting started example

9089270

Add generated Pluto notebooks to gitignore

e163b15

remove old version of example

5774a21

add Pluto and PlutoStaticHTML as deps

e7c60ce

Adapt make and build files to generate Pluto notebooks statically and…

ee9e1ca

… then render into docs

Merge branch 'main' into add-doc-examples

fd5d554

include build step for rendering Pluto notebooks

359a3c2

remove old version of the getting started example

ca2707e

update

f7816de

added doc deps

0b4bfe2

started getting started notebook

36fd0ea

Update getting_started.jl

58e9744

Getting started example

6513caf

Add generated Pluto notebooks to gitignore

1c740db

remove old version of example

a37c77a

add Pluto and PlutoStaticHTML as deps

66782b8

Adapt make and build files to generate Pluto notebooks statically and…

6853614

… then render into docs

include build step for rendering Pluto notebooks

2fbe618

remove old version of the getting started example

d97c92a

Merge branch 'main' into add-doc-examples

6019674

Update EpiAware/docs/src/examples/getting_started.jl

e805792

Merge branch 'add-doc-examples' of https://github.com/CDCgov/Rt-witho…

0189ac1

…ut-renewal into add-doc-examples

Remove unnecessary CSV dep in test

14b5d30

fix getting started

f825d38

Merge branch 'main' into add-pathfinder-init

0853042

pathfinder first pass

ac788d0

seabbs reviewed Mar 7, 2024

View reviewed changes

EpiAware/docs/src/examples/getting_started.jl Show resolved Hide resolved

change to generated_y_t approach

40f3a5b

seabbs reviewed Mar 7, 2024

View reviewed changes

EpiAware/docs/src/examples/getting_started.jl Outdated Show resolved Hide resolved

seabbs reviewed Mar 7, 2024

View reviewed changes

EpiAware/docs/src/examples/getting_started.jl Outdated Show resolved Hide resolved

seabbs reviewed Mar 7, 2024

View reviewed changes

EpiAware/docs/src/examples/getting_started.jl Show resolved Hide resolved

Update EpiAware/docs/src/examples/getting_started.jl

9b2b509

Co-authored-by: Sam Abbott <[email protected]>

SamuelBrand1 mentioned this pull request Mar 7, 2024

Passing initialisations to NUTS/inference approach #120

Closed

seabbs requested changes Mar 7, 2024

View reviewed changes

SamuelBrand1 added 2 commits March 7, 2024 11:56

fix rw_init

a8ca3e8

Merge branch 'add-pathfinder-init' of https://github.com/CDCgov/Rt-wi…

957944a

…thout-renewal into add-pathfinder-init

SamuelBrand1 added 3 commits March 7, 2024 12:16

Add warning about fail of all initial pathfinder runs

42b0b1b

Increased testing cases for manypathfinder

902d206

Make total fail a verbose exception error

67427fb

SamuelBrand1 mentioned this pull request Mar 7, 2024

Common resources for testing #121

Closed

SamuelBrand1 requested a review from seabbs March 7, 2024 12:56

seabbs reviewed Mar 7, 2024

View reviewed changes

Change manypathfinder to have some API as multipathfinder

6deb30a

SamuelBrand1 mentioned this pull request Mar 8, 2024

Clarity on use of Doc extension tools #126

Closed

SamuelBrand1 requested a review from seabbs March 8, 2024 11:55

seabbs approved these changes Mar 8, 2024

View reviewed changes

seabbs merged commit 4e035ce into main Mar 8, 2024
10 checks passed

seabbs deleted the add-pathfinder-init branch March 8, 2024 12:30

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Issue 56: Add pathfinder init #117

Issue 56: Add pathfinder init #117

SamuelBrand1 commented Mar 6, 2024 •

edited

Loading

SamuelBrand1 commented Mar 7, 2024

seabbs left a comment

seabbs commented Mar 7, 2024

codecov-commenter commented Mar 7, 2024 •

edited

Loading

SamuelBrand1 commented Mar 7, 2024

seabbs left a comment

SamuelBrand1 commented Mar 8, 2024

SamuelBrand1 commented Mar 8, 2024

seabbs commented Mar 8, 2024

seabbs left a comment

SamuelBrand1 commented Mar 8, 2024

	function _epi_aware(y_t,
	time_steps;
	epi_model::AbstractEpiModel,
	latent_model::AbstractLatentModel,
	observation_model::AbstractObservationModel,
	nsamples,
	nchains,
	pf_ndraws = 50,
	pf_nruns = 4,
	fixed_parameters = (;),
	pos_shift = 1e-6,
	executor = Transducers.ThreadedEx(),
	adtype = AutoReverseDiff(true),
	maxiters = 10,
	kwargs...)
	gen_mdl = make_epi_aware(missing, time_horizon; epi_model,
	latent_model, observation_model, pos_shift) \|>
	mdl -> fix(mdl, fixed_parameters)

	mdl = make_epi_aware(y_t, time_horizon; epi_model,
	latent_model, observation_model, pos_shift) \|>
	mdl -> fix(mdl, fixed_parameters)

	safe_mdl = make_epi_aware(y_t, time_horizon, Val(:safe_mode); epi_model,
	latent_model, observation_model, pos_shift) \|>
	mdl -> fix(mdl, fixed_parameters)

	mpf = multipathfinder(safe_mdl, max(pf_ndraws, nchains);
	nruns = pf_nruns,
	executor,
	maxiters,
	kwargs...)

	init_params = collect.(eachrow(mpf.draws_transformed.value[1:nchains, :, 1]))

	chn = sample(mdl,
	NUTS(; adtype),
	MCMCThreads(),
	nsamples ÷ nchains,
	nchains;
	init_params = init_params,
	drop_warmup = true)

	return chn, (; pathfinder_res = mpf,
	inference_mdl = mdl,
	generative_mdl = gen_mdl)
	end

Issue 56: Add pathfinder init #117

Issue 56: Add pathfinder init #117

Conversation

SamuelBrand1 commented Mar 6, 2024 • edited Loading

SamuelBrand1 commented Mar 7, 2024

seabbs left a comment

Choose a reason for hiding this comment

seabbs commented Mar 7, 2024

codecov-commenter commented Mar 7, 2024 • edited Loading

Codecov Report

SamuelBrand1 commented Mar 7, 2024

seabbs left a comment

Choose a reason for hiding this comment

SamuelBrand1 commented Mar 8, 2024

SamuelBrand1 commented Mar 8, 2024

seabbs commented Mar 8, 2024

seabbs left a comment

Choose a reason for hiding this comment

SamuelBrand1 commented Mar 8, 2024

SamuelBrand1 commented Mar 6, 2024 •

edited

Loading

codecov-commenter commented Mar 7, 2024 •

edited

Loading